Skip to content

[AutoDiff] Fix differentiable_function-related specialization crashes.#29800

Merged
dan-zheng merged 1 commit intotensorflowfrom
autodiff-generic-specialization-crash-fix
Feb 13, 2020
Merged

[AutoDiff] Fix differentiable_function-related specialization crashes.#29800
dan-zheng merged 1 commit intotensorflowfrom
autodiff-generic-specialization-crash-fix

Conversation

@dan-zheng
Copy link
Contributor

@dan-zheng dan-zheng commented Feb 12, 2020

Fix crashes related to generic specialization of partial_apply operands to
differentiable_function instructions.

differentiable_function requires derivative function operand types to match
expected derivative function types computed from the original function operand's
type, so operands cannot be specialized individually without specializing the
others.

Resolves TF-891 and TF-1126.

Fixed verification error from TF-1126:

Looking for a function: $ss4SIMDPss14DifferentiableRzSB6Scalars11SIMDStoragePRpzsAA13TangentVectorsACPRpzSBAhI_AdFRPzrlE12_vjpSubtract3lhs3rhsx5value_AJ_AJtAJc8pullbacktx_xtFZs5SIMD8VySfG_Tg5
Expected type: @convention(method) (@in_guaranteed SIMD8<Float>, @in_guaranteed SIMD8<Float>, @thick SIMD8<Float>.Type) -> (@out SIMD8<Float>, @owned @callee_guaranteed (@in_guaranteed SIMD8<Float>) -> (@out SIMD8<Float>, @out SIMD8<Float>))
Found    type: @convention(method) (SIMD8<Float>, SIMD8<Float>, @thick SIMD8<Float>.Type) -> (@out SIMD8<Float>, @owned @callee_guaranteed (@in_guaranteed SIMD8<Float>) -> (@out SIMD8<Float>, @out SIMD8<Float>))
Assertion failed: (ReInfo.getSpecializedType() == SpecializedF->getLoweredFunctionType() && "Previously specialized function does not match expected type."), function lookupSpecialization, file /Users/swiftninjas/s4tf/swift/lib/SILOptimizer/Utils/Generics.cpp, line 1827.

@dan-zheng
[AutoDiff] Fix differentiable_function-related specialization crashes.
2ee5bdf 
Fix crashes related to generic specialization of `partial_apply` operands to
`differentiable_function` instructions.

`differentiable_function` requires derivative function operand types to match
expected derivative function types computed from the original function operand's
type, so operands cannot be specialized individually without specializing the
others.

Resolves TF-891 and TF-1126.
@dan-zheng dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label Feb 12, 2020
@dan-zheng dan-zheng requested review from marcrasi and rxwei February 12, 2020 23:29
@dan-zheng
Copy link
Contributor Author

dan-zheng commented Feb 12, 2020

@swift-ci Please test tensorflow

marcrasi
marcrasi approved these changes Feb 13, 2020
View reviewed changes
Copy link

@marcrasi marcrasi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ooo nice.

Seems like we'll want to allow specialization of these things eventually because specialization can make things much faster?

rxwei
rxwei reviewed Feb 13, 2020
View reviewed changes
Copy link
Contributor

@rxwei rxwei left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think eventually we'll want specialization to apply to the @differentiable function, not operands to the differentiable_function instruction.

@dan-zheng
Copy link
Contributor Author

dan-zheng commented Feb 13, 2020

Seems like we'll want to allow specialization of these things eventually because specialization can make things much faster?

Yes, allowing specialization if possible would be ideal!

A nice approach might be to relax differentiable_function/differentiable_function_extract verification while preserving global derivative function type consistency, à la #27923.

This avoids requiring that all function-type-rewriting transformations (e.g. generic specialization, loadable-by-address) must update original/derivative function types consistently.

@dan-zheng dan-zheng merged commit 5f62183 into tensorflow Feb 13, 2020
@dan-zheng dan-zheng deleted the autodiff-generic-specialization-crash-fix branch February 13, 2020 02:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@rxwei rxwei rxwei approved these changes

+1 more reviewer

@marcrasi marcrasi marcrasi approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

tensorflow This is for "tensorflow" branch PRs.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@dan-zheng @rxwei @marcrasi